#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the version 3 of the GNU Lesser General Public License
#   as published by the Free Software Foundation.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Copyright (c) NEC Deutschland GmbH, NEC HPC Europe
#
# $Id$

import logging
from cStringIO import StringIO

from aggmon.importer import Importer, ImporterComponent
from basic_types.metric import Metric
from urlparse import urlparse
try:
    from paramiko import SSHClient, AutoAddPolicy
except:
    print "'paramiko' module not found, 'ssh' schema won't work for for torque accounting log importer!"


# Constants
TORQUE_ACCOUNTING_PATH = "/var/spool/torque/server_priv/accounting"
TORQUE_LOG_MARKER = { "A" : "aborted",         # job has been aborted by server
                      "C" : "checkpointed",    # job has been checkpointed and stopped
                      "D" : "deleted",         # job has been deleted
                      "E" : "exited",          # job has exited (either successfully or unsuccessfully)
                      "Q" : "queued",          # job has been submitted/queued
                      "R" : "rerun",           # an attempt to rerun the job has been made
                      "S" : "started",         # an attempt to start the job has been made
                      "T" : "restarted",       # an attempt to restart the job has been made (after checkpointed)
                    }
TAIL_CMD = "cd %s; tag=$(cat tag);rm -f tag; \
while read fn; do \
  while read line; do \
    fp=$(md5sum <<<$line); \
    [ ! -f tag ] && echo -n \"$fp\" > tag; \
    [ \"$fp\" == \"$tag\" ] && break 2; \
    echo \"$line\"; \
  done < <(tac $fn); \
done < <(ls -1 [0-9]* | sort -rn)"

__all__ = ["TorqueAccountingLogImporter", "TorqueAccountingLogImporterComponent"]


# Helper function
def parseURL( url ):
    if not "://" in url:
        url = "telnet://" + url
    parsed_url = urlparse( url )
    user_name = parsed_url.username
    password = parsed_url.password
    if password == "":
        password = None
    host_name = parsed_url.hostname
    port = parsed_url.port
    if port == 0 or port is None:
        port = 22   # expect this is ssh
    path = parsed_url.path
    scheme = parsed_url.scheme
    if scheme is not None and path is not None:
        return [scheme, user_name, password, host_name, path, port]
    else:
        raise ValueError( "invalid URL: %s" % url )


class TorqueAccountingLogImporter(Importer):
    def __init__( self, url="", **kwds ):
        scheme, user_name, password, host_name, path, port = parseURL( url )
        Importer.__init__( self, **kwds )
        self._ssh_client = None
        self._scheme = scheme
        self._user_name = user_name
        self._password = password
        self._host_name = host_name or "localhost"
        self._path = path or TORQUE_ACCOUNTING_PATH
        self._port = port or 22

    def _parseTorqueAccountingLog( self, _buffer ):
        import time
        source = "Torque"
        host_name = self._host_name
        for line in reversed( StringIO( _buffer ).readlines() ):
            attributes = {}
            line = line.strip()
            try:
                # example header: 08/08/2014 11:32:23;S;50438;<payload>
                comp = line.split( ";", 3 )
                assert comp[1] in TORQUE_LOG_MARKER.keys(), "Problem with Torque accounting log (unknown event): '%s'" % line
                timestamp = int( time.mktime( time.strptime( comp[0], "%m/%d/%Y %H:%M:%S" ) ) )    # time  = log (not job) timestamp
                value = TORQUE_LOG_MARKER[comp[1]]                                                  # value = event type
                name = comp[2].replace( ".", "_" )                                                  # name  = jobid
                payload = comp[3]                                                                   # attribute set
            except:
                logging.warn( "Problem with Torque accounting log (parsing header): '%s'" % line )
                continue
            for attr_str in payload.split( " " ):
                # example payload: user=myuname group=sthaber jobname=STDIN queue=workq ctime=1407490338 qtime=1407490338 etime=1407490338 start=1407490343
                #     owner=myuname@sb-master exec_host=<exec_host> Resource_List.ncpus=1 Resource_List.neednodes=1:ppn=1:foo:ghz-2.6:mhz-2601:ddr1866:qlogic
                #     Resource_List.nodect=1 Resource_List.nodes=1:ppn=1:foo:ghz-2.6:mhz-2601:ddr1866:qlogic
                attr_str = attr_str.replace( ".", "_" )
                try:
                    k, v = attr_str.split( "=", 1 )
                except:
                    logging.warn( "Problem with Torque accounting log (attribute): '%s'" % line )
                    break
                if k in ["ctime", "qtime", "etime", "start", "end"]:
                    attributes[k] = int( v )
                if k in ["user", "group", "jobname", "queue", "owner"]:
                    attributes[k] = v
                if k == "exec_host":
                    # exec_host example: sabi36/19+sabi36/18+sabi36/17+sabi36/16+sabi36/15+sabi36/14+sabi36/13+sabi36/12+sabi36/11...
                    attributes["cnodes"] = list( set( [r.split( "/" )[0] for r in v.split( "+" )] ) )
                if k.startswith( "Resource_List" ):
                    k = k.replace( "Resource_List", "RL" )
                    attributes[k] = v
            metric = Metric( name, source, time=timestamp, host=host_name, value=value, **attributes )
            yield metric

    # opens ssh connection and retrieves lines through TAIL_CMD
    def get_lines_by_ssh( self ):
        if self._ssh_client is None:
            client = SSHClient()
            client.load_system_host_keys()
            client.set_missing_host_key_policy( AutoAddPolicy() )
            self._ssh_client = client
        if self._ssh_client.get_transport() is None or \
            not self._ssh_client.get_transport().is_authenticated():
            try:
                self._ssh_client.connect( self._host_name, port=self._port,
                                          username=self._user_name, password=self._password, timeout=15 )
            except Exception, e:
                logging.error( "ssh connection to %s failed" % self._host_name )
                raise e
            self._ssh_client.get_transport().set_keepalive( 10 )
        _stdin, stdout, _stderr = self._ssh_client.exec_command( TAIL_CMD % self._path )
        _buffer = stdout.read()
        return _buffer

    def get_local_lines( self ):
        import subprocess
        p = subprocess.Popen( TAIL_CMD % self._path, shell="/bin/bash", stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, _stderr = p.communicate()
        return stdout

    def poll( self ):
        if self._scheme == "ssh":
            _buffer = self.get_lines_by_ssh()
        elif self._scheme == "file":
            _buffer = self.get_local_lines()
        else:
            raise ValueError( "Unsupported scheme: %s" % self._scheme )
        for metric in self._parseTorqueAccountingLog( _buffer ):
            yield metric


class TorqueAccountingLogImporterComponent(ImporterComponent):
    def __init__( self, *args, **kwds ):
        ImporterComponent.__init__( self, *args, **kwds )
        self.allowed_kwds.extend( ["url"] )

    def instantiate( self, *args, **kwds ):
        self.instance = None
        try:
            logging.debug( "instantiating: %s" % self.name )
            self.instance = TorqueAccountingLogImporter( *args, **kwds )
        except Exception, e:
            raise e
